#!/usr/bin/env python3
"""
数据管理器 - 基于真实数据的RAG系统
"""

import pandas as pd
import json
import os
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class Resume:
    """简历数据模型"""
    id: str
    content: str
    html_content: str = ""
    category: str = ""
    skills: List[str] = None
    experience_years: Optional[int] = None
    location: str = ""
    name: str = ""
    
    def __post_init__(self):
        if self.skills is None:
            self.skills = []

@dataclass  
class Job:
    """岗位数据模型"""
    id: str
    title: str
    short_title: str = ""
    location: str = ""
    schedule_type: str = ""
    work_from_home: bool = False
    country: str = ""
    salary_year_avg: Optional[int] = None
    company_name: str = ""
    job_skills: List[str] = None
    required_skills: List[str] = None
    preferred_skills: List[str] = None
    min_experience_years: int = 0
    description: str = ""
    
    def __post_init__(self):
        if self.job_skills is None:
            self.job_skills = []
        if self.required_skills is None:
            self.required_skills = []
        if self.preferred_skills is None:
            self.preferred_skills = []

class DataManager:
    """数据管理器 - 处理简历和岗位数据"""
    
    def __init__(self, data_dir: str = "/home/cwx/workspace/BoLe/data"):
        self.data_dir = data_dir
        self.resumes: List[Resume] = []
        self.jobs: List[Job] = []
        self.resume_df = None
        self.job_df = None
        
    def load_resume_data(self) -> bool:
        """加载简历数据"""
        try:
            resume_path = os.path.join(self.data_dir, "resume_dataset/Resume/Resume.csv")
            if not os.path.exists(resume_path):
                logger.error(f"简历数据文件不存在: {resume_path}")
                return False
            
            logger.info(f"正在加载简历数据: {resume_path}")
            self.resume_df = pd.read_csv(resume_path)
            
            logger.info(f"简历数据加载完成，共 {len(self.resume_df)} 条记录")
            
            # 转换为Resume对象
            self.resumes = []
            for idx, row in self.resume_df.head(1000).iterrows():  # 限制数量避免内存问题
                resume = Resume(
                    id=str(row.get('ID', idx)),
                    content=str(row.get('Resume_str', '')),
                    html_content=str(row.get('Resume_html', '')),
                    category=str(row.get('Category', ''))
                )
                
                # 从简历内容中提取技能和信息
                self._extract_resume_info(resume)
                self.resumes.append(resume)
            
            logger.info(f"成功处理 {len(self.resumes)} 个简历对象")
            return True
            
        except Exception as e:
            logger.error(f"加载简历数据失败: {e}")
            return False
    
    def load_job_data(self) -> bool:
        """加载岗位数据"""
        try:
            job_path = os.path.join(self.data_dir, "job_dataset/job_data_sample.csv")
            if not os.path.exists(job_path):
                logger.error(f"岗位数据文件不存在: {job_path}")
                return False
            
            logger.info(f"正在加载岗位数据: {job_path}")
            self.job_df = pd.read_csv(job_path)
            
            logger.info(f"岗位数据加载完成，共 {len(self.job_df)} 条记录")
            
            # 转换为Job对象
            self.jobs = []
            for idx, row in self.job_df.iterrows():
                # 解析技能字符串
                job_skills = self._parse_skills_string(str(row.get('job_skills', '[]')))
                required_skills = self._parse_skills_string(str(row.get('required_skills', '[]')))
                preferred_skills = self._parse_skills_string(str(row.get('preferred_skills', '[]')))
                
                job = Job(
                    id=str(idx),
                    title=str(row.get('job_title', '')),
                    short_title=str(row.get('job_title_short', '')),
                    location=str(row.get('job_location', '')),
                    schedule_type=str(row.get('job_schedule_type', '')),
                    work_from_home=bool(row.get('job_work_from_home', False)),
                    country=str(row.get('job_country', '')),
                    salary_year_avg=self._safe_int(row.get('salary_year_avg')),
                    company_name=str(row.get('company_name', '')),
                    job_skills=job_skills,
                    required_skills=required_skills,
                    preferred_skills=preferred_skills,
                    min_experience_years=self._safe_int(row.get('min_experience_years', 0)),
                    description=str(row.get('job_description', ''))
                )
                self.jobs.append(job)
            
            logger.info(f"成功处理 {len(self.jobs)} 个岗位对象")
            return True
            
        except Exception as e:
            logger.error(f"加载岗位数据失败: {e}")
            return False
    
    def _parse_skills_string(self, skills_str: str) -> List[str]:
        """解析技能字符串"""
        try:
            if skills_str in ['nan', '', 'NaN']:
                return []
            
            # 尝试解析JSON格式的技能列表
            if skills_str.startswith('[') and skills_str.endswith(']'):
                skills = json.loads(skills_str.replace("'", '"'))
                return [str(skill).strip() for skill in skills if skill]
            
            # 如果不是JSON格式，按逗号分割
            return [skill.strip() for skill in skills_str.split(',') if skill.strip()]
            
        except Exception as e:
            logger.warning(f"解析技能字符串失败: {skills_str}, 错误: {e}")
            return []
    
    def _safe_int(self, value) -> int:
        """安全转换为整数"""
        try:
            if pd.isna(value):
                return 0
            return int(float(value))
        except:
            return 0
    
    def _extract_resume_info(self, resume: Resume):
        """从简历内容中提取信息"""
        content = resume.content.lower()
        
        # 提取技能关键词（简单的关键词匹配）
        skill_keywords = [
            'python', 'java', 'javascript', 'react', 'vue', 'nodejs', 'angular',
            'sql', 'mysql', 'postgresql', 'mongodb', 'redis',
            'aws', 'docker', 'kubernetes', 'git', 'linux',
            'machine learning', 'data science', 'ai', 'tensorflow', 'pytorch',
            'project management', 'team lead', 'management'
        ]
        
        found_skills = []
        for skill in skill_keywords:
            if skill in content:
                found_skills.append(skill.title())
        
        resume.skills = found_skills
        
        # 尝试提取经验年数（简单的正则匹配）
        import re
        exp_patterns = [
            r'(\d+)\+?\s*years?\s*(?:of\s*)?(?:experience|work)',
            r'(\d+)\+?\s*years?\s*in',
            r'(\d+)\+?\s*year'
        ]
        
        for pattern in exp_patterns:
            match = re.search(pattern, content)
            if match:
                try:
                    resume.experience_years = int(match.group(1))
                    break
                except:
                    continue
        
        # 提取姓名（从简历开头提取，简单处理）
        lines = resume.content.split('\n')[:10]  # 前10行
        for line in lines:
            line = line.strip()
            if line and len(line) < 50 and not any(keyword in line.lower() for keyword in ['summary', 'objective', 'experience', 'education']):
                # 可能是姓名
                if len(line.split()) <= 3:  # 姓名通常不超过3个词
                    resume.name = line
                    break
    
    def search_resumes(self, 
                      keywords: List[str] = None, 
                      location: str = None, 
                      min_experience: int = None,
                      category: str = None,
                      limit: int = 20) -> List[Resume]:
        """搜索简历"""
        results = self.resumes.copy()
        
        # 关键词搜索
        if keywords:
            filtered_results = []
            for resume in results:
                content_lower = resume.content.lower()
                if any(keyword.lower() in content_lower for keyword in keywords):
                    filtered_results.append(resume)
            results = filtered_results
        
        # 经验年限筛选
        if min_experience is not None:
            results = [r for r in results if r.experience_years and r.experience_years >= min_experience]
        
        # 类别筛选
        if category:
            results = [r for r in results if category.lower() in r.category.lower()]
        
        return results[:limit]
    
    def search_jobs(self, 
                   keywords: List[str] = None,
                   location: str = None,
                   min_salary: int = None,
                   company: str = None) -> List[Job]:
        """搜索岗位"""
        results = self.jobs.copy()
        
        # 关键词搜索
        if keywords:
            filtered_results = []
            for job in results:
                text_fields = [job.title, job.description] + job.job_skills + job.required_skills
                text_content = ' '.join(text_fields).lower()
                if any(keyword.lower() in text_content for keyword in keywords):
                    filtered_results.append(job)
            results = filtered_results
        
        # 位置筛选
        if location:
            results = [j for j in results if location.lower() in j.location.lower()]
        
        # 薪资筛选
        if min_salary:
            results = [j for j in results if j.salary_year_avg and j.salary_year_avg >= min_salary]
        
        # 公司筛选
        if company:
            results = [j for j in results if company.lower() in j.company_name.lower()]
        
        return results
    
    def get_resume_by_id(self, resume_id: str) -> Optional[Resume]:
        """根据ID获取简历"""
        for resume in self.resumes:
            if resume.id == resume_id:
                return resume
        return None
    
    def get_job_by_id(self, job_id: str) -> Optional[Job]:
        """根据ID获取岗位"""
        for job in self.jobs:
            if job.id == job_id:
                return job
        return None
    
    def initialize(self) -> bool:
        """初始化数据管理器"""
        logger.info("初始化数据管理器...")
        
        success = True
        if not self.load_resume_data():
            logger.error("简历数据加载失败")
            success = False
        
        if not self.load_job_data():
            logger.error("岗位数据加载失败")
            success = False
        
        if success:
            logger.info("数据管理器初始化成功")
        else:
            logger.error("数据管理器初始化失败")
        
        return success

# 全局数据管理器实例
data_manager = DataManager()

if __name__ == "__main__":
    # 测试代码
    if data_manager.initialize():
        print(f"✅ 数据加载成功")
        print(f"📄 简历数量: {len(data_manager.resumes)}")
        print(f"💼 岗位数量: {len(data_manager.jobs)}")
        
        # 测试搜索
        python_resumes = data_manager.search_resumes(keywords=['python'], limit=5)
        print(f"🔍 Python相关简历: {len(python_resumes)}个")
        
        if python_resumes:
            print(f"📋 示例简历: {python_resumes[0].name or '未知姓名'}")
    else:
        print("❌ 数据加载失败")









